import os
import shutil
import json
import yaml

from Entitys.tasks import TrainTask
from PredictTools.runScirpt import excuteCommand


def do_train(trainTask: TrainTask, train_script_path: str):
    imageTs = trainTask.imgPaths
    labelTs = trainTask.labelPaths

    file_datasets = os.path.join(trainTask.savePath, 'datas', 'raw_data')
    file_train = os.path.join(file_datasets, 'imageTr')
    file_label = os.path.join(file_datasets, 'labelTr')

    batch_copy_file(imageTs, file_train)
    batch_copy_file(labelTs, file_label)

    make_json(trainTask, file_datasets)

    config_save_path = os.path.join(trainTask.savePath, 'config.yml')
    make_train_config(config_save_path, file_datasets)
    modle_save_path = os.path.join(trainTask.savePath, 'models')

    train_cmd = f'python {train_script_path} --config {config_save_path} --log_iters 20 --precision fp16 --nnunet --save_dir {modle_save_path} --save_interval 2000 '

    print(train_cmd)

    train_log = excuteCommand(train_cmd)
    return train_log


def batch_copy_file(paths, out_path):
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    for path in paths:
        out_tmp = os.path.join(out_path, os.path.basename(path))
        shutil.copy(path, out_tmp)


def make_json(trainTask: TrainTask, file_datasets):
    labels = trainTask.labelClass
    label_dict = {}
    for i, label in enumerate(labels):
        label_dict[str(i)] = label

    training_list = []
    for i in range(len(trainTask.imgPaths)):
        training_list.append({
            'image': './' + os.path.basename(trainTask.imgPaths[i]),
            'label': './' + os.path.basename(trainTask.labelPaths[i])
        })

    train_json = {
        'root': {
            'name': 'train',
            'author': 'sfl',
            'reference': '',
            "licence": "CC-BY-SA 4.0",
            "release": "1.0 01/05/2022",
            "tensorImageSize": '3D',
            'modality': {
                '0': trainTask.image_type
            },
            'label': label_dict,
            'numTraining': len(trainTask.imgPaths),
            'numTest': 0,
            'training': training_list,
            'test': []
        }
    }

    with open(os.path.join(file_datasets, 'dataset.json'), 'w') as f:
        json.dump(train_json, f, indent=4)


def make_train_config(config_save_path, file_datasets):
    config = {
        'data_root': file_datasets,
        'batch_size': 1,
        'iters': 100000,
        'model':
            {
                'type': 'NNUNet',
                'plan_path': os.path.join(file_datasets, 'preprocessed', 'nnUNetPlansv2.1_plans_3D.pkl'),
                'stage': 0,
                'cascade': True
            },
        'train_dataset':
            {
                "type": "MSDDataset",
                "plans_name": "nnUNetPlansv2.1_plans_3D.pkl",
                "dataset_root": "/",
                "result_dir": "/",
                "raw_data_dir": os.path.join(file_datasets, 'raw_data'),
                "decathlon_dir": os.path.join(file_datasets, 'decathlon'),
                "cropped_data_dir": os.path.join(file_datasets, 'cropped'),
                "preprocessed_dir": os.path.join(file_datasets, 'preprocessed'),
                "plan2d": False,
                "plan3d": True,
                "num_batches_per_epoch": 250,
                "fold": 0,
                "stage": 0,
                "unpack_data": True,
                "cascade": True,
                "mode": "train",
            },
        'optimizer':
            {
                "type": "sgd",
                "momentum": 0.99,
                "weight_decay": 0.00003,
                "use_nesterov": True
            },
        'lr_scheduler':
            {
                "type": "PolynomialDecay",
                "learning_rate": 0.01,
                "end_lr": 0,
                "power": 0.9
            },
        'loss':
            {
                'types':
                    [
                        {
                            'type': "MultipleLoss",
                            'plan_path': os.path.join(file_datasets, 'preprocessed', 'nnUNetPlansv2.1_plans_3D.pkl'),
                            "stage": 0,
                            "losses":
                                [
                                    {
                                        'type': 'DC_and_CE_loss',
                                        'do_bg': False,
                                        'batch_dice': False,
                                    }
                                ],
                            'coef': [1.0]
                        }
                    ]
                ,
                'coef': [1.0]
            }
    }

    with open(config_save_path, 'w') as f:
        yaml.dump(config, f, indent=4)
